package com.omni.monitor.plugin.trace_monitor;

import com.omni.monitor.model.OmniConst;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.UUID;
/**
 * 基于Spring Boot的traceId拦截器
 * 必须添加@Component注解，让Spring Boot自动扫描并管理
 */
//@Component
public class TraceHttpInterceptor  implements HandlerInterceptor {
    private final TraceContext traceContext;
    // 构造注入TraceContext
    public TraceHttpInterceptor(TraceContext traceContext) {
        this.traceContext = traceContext;
    }

    /**
     * 请求处理前执行：注入traceId到MDC
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        // 1. 优先从请求头获取（分布式系统中，上游服务可能已传递）
        String traceId = request.getHeader(OmniConst.TRACE_ID_HEADER);
        String spanId = request.getHeader(OmniConst.SPAN_ID_HEADER);
        String parentSpanId = request.getHeader(OmniConst.PARENT_SPAN_ID_HEADER);

        // 2. 若请求头无，则生成新的traceId（UUID去除横线，缩短长度）
        if (StringUtils.isEmpty(traceId)) {
            traceId = UUID.randomUUID().toString().replace("-", "");
            spanId = "0"; // 根节点spanId
            parentSpanId = ""; // 根节点无父节点
        }

        // 3. 存入MDC（后续日志框架可直接引用）
        MDC.put(OmniConst.TRACE_ID_MDC_KEY, traceId);
        // 3. 初始化MDC（调用链三要素）
        MDC.put(OmniConst.TRACE_ID, traceId);
        MDC.put(OmniConst.SPAN_ID, spanId);
        MDC.put(OmniConst.PARENT_SPAN_ID, parentSpanId);
        MDC.put(OmniConst.THREAD_NAME, Thread.currentThread().getName());
        // 4. 写入响应头（方便前端/下游服务获取）
        response.setHeader(OmniConst.TRACE_ID_HEADER, traceId);
        response.setHeader(OmniConst.SPAN_ID_HEADER, spanId);
        response.setHeader(OmniConst.PARENT_SPAN_ID_HEADER, parentSpanId);

        // 必须返回true，否则请求会被拦截（无法到达Controller）
        return true;
    }

    /**
     * 请求完成后执行：清理MDC，避免线程池复用导致的上下文污染
     */
    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) {
        // 注意：Spring Boot中线程池（如Tomcat线程池）会复用线程，必须清理
        String traceId = MDC.get(OmniConst.TRACE_ID_MDC_KEY);
        MDC.remove(OmniConst.TRACE_ID_MDC_KEY);
        // 或使用MDC.clear()清理所有键（根据实际需求选择）
        // MDC.clear();
        // 关键：请求结束后清理上下文（避免Tomcat线程池复用污染）
        traceContext.clear(traceId);
    }
}
